import torch
import os

from torchvision import transforms
from PIL import Image
from ResNet import resnet50
import matplotlib.pyplot as plt
 

class predict_class():
    def __init__(self,path):
        self.load_weight()
        self.predict(path)

    def load_weight(self):
        self.weights_path= '/home/h/筑工/resnet50.pth'
        self.class_names=['1', '2', '3', '4']
        if torch.cuda.is_available():  # 如果有GPU就用，没有就用CPU
            self.device = torch.device('cuda:0')
        else:
            self.device = torch.device('cpu')

    def predict(self,image_path):
            # # 权重参数路径
            # weights_path = '/home/h/筑工/data/save_weights/resnet50.pth'
            # # 预测索引对应的类别名称
            # # class_names=['Bananaquit', 'Black Skimmer', 'Black Throated Bushtiti', 'Cockatoo']
            # class_names=['1', '2', '3', '4']
            
            # # 获取GPU设备
            # if torch.cuda.is_available():  # 如果有GPU就用，没有就用CPU
            #     device = torch.device('cuda:0')
            # else:
            #     device = torch.device('cpu')
            # -------------------------------------------------- #
            #（1）数据加载
            # -------------------------------------------------- #
            # 预处理函数
            data_transform = transforms.Compose([
                # 将输入图像的尺寸变成224*224
                transforms.Resize((224,224)),
                # 数据变成tensor类型，像素值归一化，调整维度[h,w,c]==>[c,h,w]
                transforms.ToTensor(),
                # 对每个通道的像素进行标准化，给出每个通道的均值和方差
                transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))])
            
            # 读取图片
            frame = Image.open(image_path)
            # 展示图片
            # plt.imshow(frame)
            # plt.title('Black_Throated_Bushtiti')
            # plt.show()
            
            # 数据预处理
            img = data_transform(frame)
            # 给图像增加batch维度 [c,h,w]==>[b,c,h,w]
            img = torch.unsqueeze(img, dim=0)

            # -------------------------------------------------- #
            #（2）图像预测
            # -------------------------------------------------- #
            # 加载模型
            model = resnet50(num_classes=4, include_top=True)
            # 加载权重文件
            model.load_state_dict(torch.load(self.weights_path, map_location=self.device))
            # 模型切换成验证模式，dropout和bn切换形式
            model.eval()
            
            # 前向传播过程中不计算梯度
            with torch.no_grad():
                # 前向传播
                outputs = model(img)
                # 只有一张图就挤压掉batch维度
                outputs = torch.squeeze(outputs)
                # 计算图片属于4个类别的概率
                predict = torch.softmax(outputs, dim=0)
                # 得到类别索引
                predict_cla = torch.argmax(predict).numpy()
            
            # 获取最大预测类别概率
            predict_score = round(torch.max(predict).item(), 4)
            # 获取预测类别的名称
            predict_name = self.class_names[predict_cla]
            # 展示预测结果
            # plt.imshow(frame)
            # plt.title('class: '+str(predict_name)+'\n score: '+str(predict_score))
            print('class: '+str(predict_name)+'\n score: '+str(predict_score))
#             # plt.pause(0.1)








# # -------------------------------------------------- #
# #（0）参数设置
# # -------------------------------------------------- #
# # 图片文件路径
# n = 0
# total =0
# root_path = "/home/h/筑工/data/test"
# path = os.listdir(root_path)
# for j in range(4):
#     print(str(path[j]))
#     img_path = str("".join(sorted([os.path.join(root_path,str(path[j]))])))
#     # img_path = str("".join(img_path)) 
#     print(img_path)
#     img_path = sorted([os.path.join(img_path,f) for f in os.listdir(img_path) if f.endswith(".jpg")])
#     for i in range(len(img_path)):
#         print(img_path[i])
# #         # img_path = "/home/h/筑工/data/val/3/frame_20230315-162959.jpg"
#         # 权重参数路径
#         weights_path = '/home/h/筑工/data/save_weights/99_resnet50.pth'
#         # 预测索引对应的类别名称
#         # class_names=['Bananaquit', 'Black Skimmer', 'Black Throated Bushtiti', 'Cockatoo']
#         class_names=['1', '2', '3', '4']
        
#         # 获取GPU设备
#         if torch.cuda.is_available():  # 如果有GPU就用，没有就用CPU
#             device = torch.device('cuda:0')
#         else:
#             device = torch.device('cpu')

#         # -------------------------------------------------- #
#         #（1）数据加载
#         # -------------------------------------------------- #
#         # 预处理函数
#         data_transform = transforms.Compose([
#             # 将输入图像的尺寸变成224*224
#             transforms.Resize((224,224)),
#             # 数据变成tensor类型，像素值归一化，调整维度[h,w,c]==>[c,h,w]
#             transforms.ToTensor(),
#             # 对每个通道的像素进行标准化，给出每个通道的均值和方差
#             transforms.Normalize(mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5))])
        
#         # 读取图片
#         frame = Image.open(img_path[i])
#         # 展示图片
#         # plt.imshow(frame)
#         # plt.title('Black_Throated_Bushtiti')
#         # plt.show()
        
#         # 数据预处理
#         img = data_transform(frame)
#         # 给图像增加batch维度 [c,h,w]==>[b,c,h,w]
#         img = torch.unsqueeze(img, dim=0)

#         # -------------------------------------------------- #
#         #（2）图像预测
#         # -------------------------------------------------- #
#         # 加载模型
#         model = resnet50(num_classes=4, include_top=True)
#         # 加载权重文件
#         model.load_state_dict(torch.load(weights_path, map_location=device))
#         # 模型切换成验证模式，dropout和bn切换形式
#         model.eval()
        
#         # 前向传播过程中不计算梯度
#         with torch.no_grad():
#             # 前向传播
#             outputs = model(img)
#             # 只有一张图就挤压掉batch维度
#             outputs = torch.squeeze(outputs)
#             # 计算图片属于4个类别的概率
#             predict = torch.softmax(outputs, dim=0)
#             # 得到类别索引
#             predict_cla = torch.argmax(predict).numpy()
        
#         # 获取最大预测类别概率
#         predict_score = round(torch.max(predict).item(), 4)
#         # 获取预测类别的名称
#         predict_name = class_names[predict_cla]
#         print(type(predict_name))
        
#         # 展示预测结果
#         # plt.imshow(frame)
#         # plt.title('class: '+str(predict_name)+'\n score: '+str(predict_score))
#         if (predict_name!=str(path[j])):
#             n=n+1
#         print(predict_name,str(path[j]))
#         # plt.pause(0.1)
#     total = len(img_path)+total
# print(n,total,n/total)
        
        # plt.show()